/*
 * @Author: Wangjun
 * @Date: 2021-05-15 22:43:20
 * @LastEditTime: 2025-04-03 09:57:54
 * @LastEditors: wangjun haodreams@163.com
 * @Description:
 * @FilePath: \xr_node_calcd:\go\src\gitee.com\haodreams\libs\wrap\wrap.go
 * hnxr
 */
package wrap

import (
	"sync"
)

// WaitGroupWrap 支持设置最大线程数的wg
func WaitGroupWrap[G any](maxThread int, chanAny chan *G, cb func(g *G)) *sync.WaitGroup {
	wg := new(sync.WaitGroup)
	wg.Add(maxThread)
	for i := 0; i < maxThread; i++ {
		go func() {
			for g := range chanAny {
				cb(g)
			}
			wg.Done()
		}()
	}
	return wg
}

// Deprecated WaitGroupWrapper 建议使用WaitGroupWrap
type WaitGroupWrapper struct {
	sync.WaitGroup
	current  int32
	maxCount int32 //最大线程数
	cond     *sync.Cond
	lock     sync.Mutex
}

// Deprecated NewWaitGroupWrapper 支持设置最大线程数的wg
func NewWaitGroupWrapper(maxThread int) *WaitGroupWrapper {
	wg := new(WaitGroupWrapper)
	wg.maxCount = int32(maxThread)
	wg.cond = sync.NewCond(&wg.lock)
	return wg
}

// Acquire 尝试获取资源，若没有资源则等待
func (m *WaitGroupWrapper) Acquire() {
	m.lock.Lock()
	for m.current >= m.maxCount {
		m.cond.Wait()
	}
	m.current++
	m.lock.Unlock()
}

// Release 释放资源
func (m *WaitGroupWrapper) Release() {
	m.lock.Lock()
	m.current--
	m.cond.Signal()
	m.lock.Unlock()
}

// Wrap cb 回调函数
// 回调函数需要注意的是 回调函数内部的变量是引用变量，需要注意创建新变量复制
func (m *WaitGroupWrapper) Wrap(cb func()) {
	if m.maxCount == 0 {
		cb()
		return
	}
	m.Acquire()
	m.Add(1)
	go func() {
		cb()
		m.Release()
		m.Done()
	}()
}
